# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-2025.2.14
import abc
from typing import Any
from langchain.tools import BaseTool
import logging
import pandas as pd
import sqlite3
import utils.configFinRAG as configFinRAG
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM

# 全局变量初始化
g_example_question_list = []
g_example_sql_list = []
g_example_info_list = []
g_example_fa_list = []
g_example_token_list = []

# 配置日志
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger()

# 配置数据库路径
DATABASE_PATH = configFinRAG.database_path


def load_example_data():
    """加载示例数据到全局变量"""
    if not g_example_question_list:
        sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
        for _, row in sql_examples_file.iterrows():
            g_example_question_list.append(row['问题'])
            g_example_sql_list.append(row['SQL'])
            g_example_info_list.append(row['资料'])
            g_example_fa_list.append(row['FA'])
            tokens = TOKENIZER(row['问题'])['input_ids']
            g_example_token_list.append(tokens)
        logger.info("示例数据加载完成")


class FinSQLRAG(BaseTool, abc.ABC):
    name = "查询金融数据库"
    description = "当被问到金融查询相关的问题时，会去金融数据库检索结果"

    def __init__(self):
        super().__init__()

    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
        # 用例中没有用到 arun 不予具体实现
        pass

    def _run(self, para: str) -> str:
        query = para
        try:
            # 确保示例数据已加载
            load_example_data()

            # 生成SQL语句
            result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list, g_example_token_list)
            logger.info(f"Generated SQL: {sql}")

            # 查询数据库
            with sqlite3.connect(DATABASE_PATH) as conn:
                cursor = conn.cursor()
                success_flag, exc_result = query_db(sql, cursor)
                if not success_flag:
                    logger.error(f"Database query failed: {exc_result}")
                    return "FinSQLRAG处理异常！"

            # 生成答案
            answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list, g_example_fa_list, g_example_token_list)
            return answer

        except Exception as e:
            logger.error(f"Retriever Error: {e}")
            return "FinSQLRAG处理异常！"


if __name__ == "__main__":
    tool = FinSQLRAG()
    result = tool.run("请帮我计算，在20210715，中信行业分类划分的一级行业为消费者服务行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。")
    print(result)
